import matplotlib.pyplot as plt
import pandas as pd

# print(plt.style.available)
plt.style.use('fivethirtyeight')
train_data = pd.read_csv('./train.tsv', sep='\t')
train_data['sentence_length'] = list(map(lambda x: len(x), train_data['sentence']))

test_data = pd.read_csv('./dev.tsv', sep='\t')
test_data['sentence_length'] = list(map(lambda x: len(x), test_data['sentence']))

train_data_0 = len(train_data[train_data['label'] == 0])
train_data_1 = len(train_data[train_data['label'] == 1])
print(train_data_0, train_data_1)

train_data_bars = [train_data_0, train_data_1]
test_data_0 = len(test_data[test_data['label'] == 0])
test_data_1 = len(test_data[test_data['label'] == 1])
print(test_data_0, test_data_1)

test_data_bars = [test_data_0, test_data_1]
#
# sentence_length = [train_data['sentence_length'], test_data['sentence_length']]
# print(sentence_length)
# 绘制图像
plt.figure(figsize=(16, 16))
label_train = ['Train 0', 'Train 1']
label_test = ['Test 0', 'Test 1']
plt.subplot(221)

plt.bar(label_train, train_data_bars, color=['blue', 'orange'])
plt.title('Train Data Distribution')
plt.ylabel('Count')

plt.subplot(222)
plt.bar(label_test, test_data_bars, color=['blue', 'orange'])
plt.title('Test Data Distribution')
plt.ylabel('Count')

plt.subplot(223)
plt.hist(train_data['sentence_length'], bins=300)
plt.title('Test Data Distribution')
plt.ylabel('Count')

plt.subplot(224)
plt.hist(test_data['sentence_length'], bins=300)
plt.title('Test Data Distribution')
plt.ylabel('Count')

# 保存图像到文件（可选）
plt.savefig('softmax_output.png')

# 调整布局并显示
# plt.tight_layout()
# 显示图像
plt.show()
